## 成交量加权平均价格 VWAP（优化版本），用于判断价格相对于VWAP的位置
## 条件：当前价格 > VWAP * 倍数（动态调整，基础1.0倍）

from xtquant import xtdata
from typing import Dict, Optional, Union
import pandas as pd
from datetime import datetime
import time


class IndependentVWAPCalculator:
    """
    完全独立的VWAP计算器
    
    特性：
    1. 不依赖订阅，直接使用xtdata获取历史数据
    2. 输入股票代码即可获取最新VWAP和价格判断
    3. 自动下载当日分钟数据并计算
    4. 支持缓存，避免重复下载
    5. 动态调整倍数，根据数据量自适应
    """
    
    def __init__(self, period: int = 15, factor: float = 1.0, xtdata=None):
        if period <= 0:
            raise ValueError("VWAP period must be positive")
        if factor <= 0:
            raise ValueError("VWAP factor must be positive")
            
        self.period: int = period
        self.factor: float = factor
        self.xtdata = xtdata
        
        # 每个股票的状态
        self._price_history: Dict[str, list] = {}
        self._volume_history: Dict[str, list] = {}
        self._last_vwap: Dict[str, float] = {}
        self._last_price: Dict[str, float] = {}
        self._last_volume: Dict[str, float] = {}
        
        # 缓存机制：避免重复下载同一天的数据
        self._cache_date: Dict[str, str] = {}
        self._cache_ttl: int = 60  # 缓存有效期（秒）- 1分钟更新一次确保获取最新完整分钟数据
        self._cache_time: Dict[str, float] = {}
        self._cache_minute: Dict[str, str] = {}  # 股票代码 -> 缓存的分钟时间戳
        
        # 交易时间段常量
        self.MORNING_START = datetime.strptime("09:30", "%H:%M").time()
        self.MORNING_END = datetime.strptime("11:30", "%H:%M").time()
        self.AFTERNOON_START = datetime.strptime("13:00", "%H:%M").time()
        self.AFTERNOON_END = datetime.strptime("15:00", "%H:%M").time()
    
    def reset(self, stock_code: Optional[str] = None) -> None:
        """重置指定股票或全部股票的VWAP状态"""
        if stock_code is None:
            self._price_history.clear()
            self._volume_history.clear()
            self._last_vwap.clear()
            self._last_price.clear()
            self._last_volume.clear()
            self._cache_date.clear()
            self._cache_time.clear()
            self._cache_minute.clear()
            return
        
        self._price_history.pop(stock_code, None)
        self._volume_history.pop(stock_code, None)
        self._last_vwap.pop(stock_code, None)
        self._last_price.pop(stock_code, None)
        self._last_volume.pop(stock_code, None)
        self._cache_date.pop(stock_code, None)
        self._cache_time.pop(stock_code, None)
        self._cache_minute.pop(stock_code, None)
    
    def _is_cache_valid(self, stock_code: str, today: str) -> bool:
        """检查缓存是否有效"""
        if stock_code not in self._cache_date or stock_code not in self._cache_time:
            return False
        
        # 检查日期是否匹配
        if self._cache_date[stock_code] != today:
            return False
        
        # 检查当前分钟是否与缓存分钟相同
        current_minute = datetime.now().strftime("%Y%m%d%H%M")
        cached_minute = self._cache_minute.get(stock_code, "")
        if cached_minute != current_minute:
            return False
        
        # 检查时间是否在有效期内
        current_time = time.time()
        cache_time = self._cache_time[stock_code]
        return (current_time - cache_time) < self._cache_ttl
    
    def _download_minute_data(self, stock_code: str, date: str) -> bool:
        """下载指定股票的当日分钟数据"""
        if not self.xtdata:
            print("⚠️ xtdata未初始化，无法下载数据")
            return False
        
        try:
            self.xtdata.download_history_data(
                stock_code=stock_code, 
                period="1m", 
                start_time=date, 
                end_time=date
            )
            time.sleep(0.1)  # 避免请求过快
            return True
        except Exception as e:
            print(f"⚠️ {stock_code} 分钟数据下载失败: {e}")
            return False
    
    def _get_minute_data(self, stock_code: str, date: str) -> Optional[pd.DataFrame]:
        """获取指定股票的当日分钟数据"""
        if not self.xtdata:
            return None
        
        try:
            data = self.xtdata.get_local_data(
                stock_list=[stock_code], 
                period='1m', 
                start_time=date, 
                end_time=date
            )
            
            df = data.get(stock_code) if data else None
            if df is None or df.empty:
                return None
            
            # 识别时间列、价格列和成交量列
            time_col = next((c for c in df.columns if 'time' in c.lower()), None)
            price_col = next((c for c in ['close', 'lastPrice', 'price', 'last', 'current'] if c in df.columns), None)
            volume_col = next((c for c in ['volume', 'vol', 'amount'] if c in df.columns), None)
            
            if not time_col or not price_col or not volume_col:
                return None
            
            # 时间格式统一处理
            if str(df[time_col].dtype) in ('int64', 'float64'):
                df['__time__'] = pd.to_datetime(df[time_col], unit='ms') + pd.Timedelta(hours=8)
            else:
                try:
                    df['__time__'] = pd.to_datetime(df[time_col], format='%Y%m%d%H%M%S', errors='coerce')
                except Exception:
                    df['__time__'] = pd.to_datetime(df[time_col], errors='coerce')
            
            df = df[df['__time__'].notna()].sort_values('__time__')
            
            # 添加列名到DataFrame以便后续使用
            df._price_col = price_col
            df._volume_col = volume_col
            
            return df if not df.empty else None
            
        except Exception as e:
            print(f"⚠️ {stock_code} 分钟数据获取失败: {e}")
            return None
    
    def _filter_trading_session_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """筛选交易时段内的数据，并排除当前未完成的分钟"""
        try:
            # 先筛选交易时段内的数据
            mask = [((self.MORNING_START <= ts.time() <= self.MORNING_END) or
                     (self.AFTERNOON_START <= ts.time() <= self.AFTERNOON_END)) for ts in df['__time__']]
            if any(mask):
                filtered_df = df[mask]
                
                # 排除当前未完成的分钟：只保留完整的分钟数据
                current_time = datetime.now()
                current_minute = current_time.replace(second=0, microsecond=0)
                
                # 只保留时间小于当前分钟的数据（即已完成的分钟）
                complete_mask = filtered_df['__time__'] < current_minute
                if complete_mask.any():
                    return filtered_df[complete_mask]
                else:
                    return filtered_df  # 如果没有完整分钟数据，返回原数据
            else:
                return df
        except Exception:
            pass
        return df
    
    def _calculate_vwap_from_data(self, stock_code: str, df: pd.DataFrame) -> Optional[float]:
        """从分钟数据计算VWAP"""
        try:
            # 安全地获取价格列名和成交量列名
            price_col = getattr(df, '_price_col', None)
            volume_col = getattr(df, '_volume_col', None)
            if not price_col or not volume_col:
                # 如果没有属性，重新识别列
                price_col = next((c for c in ['close', 'lastPrice', 'price', 'last', 'current'] if c in df.columns), None)
                volume_col = next((c for c in ['volume', 'vol', 'amount'] if c in df.columns), None)
                if not price_col or not volume_col:
                    print(f"⚠️ {stock_code} 无法识别价格列或成交量列")
                    return None
            
            # 获取价格和成交量数据
            prices = df[price_col].dropna().tolist()
            volumes = df[volume_col].dropna().tolist()
            
            if len(prices) < 2 or len(volumes) < 2:
                return None
            
            # 确保价格和成交量数据长度一致
            min_len = min(len(prices), len(volumes))
            prices = prices[:min_len]
            volumes = volumes[:min_len]
            
            # 保存价格和成交量历史
            self._price_history[stock_code] = prices
            self._volume_history[stock_code] = volumes
            
            # 计算VWAP（成交量加权平均价格）
            if len(prices) >= self.period:
                # 使用指定周期的数据
                recent_prices = prices[-self.period:]
                recent_volumes = volumes[-self.period:]
            else:
                # 数据不足时，使用所有可用数据
                recent_prices = prices
                recent_volumes = volumes
            
            # VWAP = Σ(价格 × 成交量) / Σ(成交量)
            total_volume_price = sum(p * v for p, v in zip(recent_prices, recent_volumes))
            total_volume = sum(recent_volumes)
            
            if total_volume > 0:
                vwap = total_volume_price / total_volume
            else:
                # 如果成交量为0，使用简单平均价格
                vwap = sum(recent_prices) / len(recent_prices)
            
            self._last_vwap[stock_code] = vwap
            self._last_price[stock_code] = prices[-1] if prices else 0
            self._last_volume[stock_code] = volumes[-1] if volumes else 0
            
            return vwap
            
        except Exception as e:
            print(f"⚠️ {stock_code} VWAP计算异常: {e}")
            return None
    
    def _get_dynamic_factor(self, stock_code: str) -> float:
        """根据数据量获取动态倍数"""
        if stock_code not in self._price_history:
            return self.factor
        
        data_count = len(self._price_history[stock_code])
        
        # 数据量少时降低倍数要求
        if data_count < 3:
            return self.factor * 0.7
        elif data_count < 6:
            return self.factor * 0.8
        elif data_count < self.period:
            return self.factor * 0.9
        else:
            return self.factor
    
    def _filter_data_before_time(self, df: pd.DataFrame, end_time: datetime) -> pd.DataFrame:
        """
        过滤掉指定时间之后的数据（回测模式使用）
        
        Args:
            df: 原始数据DataFrame
            end_time: 结束时间
            
        Returns:
            过滤后的DataFrame
        """
        try:
            if df.empty:
                return df
            
            # 确保时间列存在
            if '__time__' not in df.columns:
                return df
            
            # 过滤掉结束时间之后的数据
            filtered_df = df[df['__time__'] <= end_time]
            return filtered_df
            
        except Exception as e:
            print(f"⚠️ 数据过滤异常: {e}")
            return df
    
    def get_latest_vwap(self, stock_code: str, end_time: Optional[datetime] = None) -> Optional[float]:
        """
        获取指定股票的最新VWAP值
        
        Args:
            stock_code: 股票代码
            end_time: 结束时间（回测模式使用，实时模式为None）
            
        Returns:
            VWAP值，如果无法计算则返回None
        """
        try:
            # 确定当前日期
            if end_time is not None:
                # 回测模式：使用指定的结束时间
                today = end_time.strftime("%Y%m%d")
                current_time = end_time
            else:
                # 实时模式：使用当前时间
                today = datetime.now().strftime("%Y%m%d")
                current_time = datetime.now()
            
            # 检查缓存是否有效
            if self._is_cache_valid(stock_code, today):
                return self._last_vwap.get(stock_code)
            
            # 缓存无效，需要重新下载和计算
            print(f"🔄 {stock_code} 开始下载当日分钟数据...")
            
            # 下载分钟数据
            if not self._download_minute_data(stock_code, today):
                return None
            
            # 获取分钟数据
            df = self._get_minute_data(stock_code, today)
            if df is None:
                return None
            
            # 筛选交易时段内的数据
            df = self._filter_trading_session_data(df)
            if df.empty:
                return None
            
            # 在回测模式下，过滤掉结束时间之后的数据
            if end_time is not None:
                df = self._filter_data_before_time(df, end_time)
                if df.empty:
                    return None
            
            # 计算VWAP
            vwap = self._calculate_vwap_from_data(stock_code, df)
            
            # 更新缓存
            if vwap is not None:
                self._cache_date[stock_code] = today
                self._cache_time[stock_code] = time.time()
                self._cache_minute[stock_code] = current_time.strftime("%Y%m%d%H%M")
            
            return vwap
            
        except Exception as e:
            print(f"❌ {stock_code} 获取VWAP异常: {e}")
            return None
    
    def check_price_vs_vwap(self, stock_code: str, end_time: Optional[datetime] = None) -> Dict[str, Union[bool, float, str]]:
        """
        检查价格相对于VWAP的位置
        
        Args:
            stock_code: 股票代码
            end_time: 结束时间（回测模式使用，实时模式为None）
            
        Returns:
            包含判断结果的字典
        """
        try:
            # 获取最新VWAP
            vwap = self.get_latest_vwap(stock_code, end_time)
            if vwap is None:
                return {
                    'is_above': False,
                    'vwap': None,
                    'current_price': None,
                    'factor': None,
                    'dynamic_factor': None,
                    'message': '无法获取VWAP数据'
                }
            
            # 获取当前价格
            current_price = self._last_price.get(stock_code, 0)
            
            # 检查是否在交易时段
            if end_time is not None:
                # 回测模式：使用指定的结束时间
                now = end_time.time()
            else:
                # 实时模式：使用当前时间
                now = datetime.now().time()
            is_trading_time = ((self.MORNING_START <= now <= self.MORNING_END) or 
                              (self.AFTERNOON_START <= now <= self.AFTERNOON_END))
            
            if current_price == 0:
                if not is_trading_time:
                    return {
                        'is_above': False,
                        'vwap': vwap,
                        'current_price': current_price,
                        'factor': self.factor,
                        'dynamic_factor': None,
                        'message': '非交易时段，使用最近交易时段价格'
                    }
                else:
                    return {
                        'is_above': False,
                        'vwap': vwap,
                        'current_price': current_price,
                        'factor': self.factor,
                        'dynamic_factor': None,
                        'message': '交易时段内价格为0'
                    }
            
            # 获取动态倍数
            dynamic_factor = self._get_dynamic_factor(stock_code)
            
            # 判断价格是否高于VWAP的倍数
            is_above = current_price > (vwap * dynamic_factor)
            
            # 生成消息
            if is_above:
                message = f"价格{current_price:.3f}>VWAP:{vwap:.3f}×{dynamic_factor:.2f}"
            else:
                message = f"价格{current_price:.3f}≤VWAP:{vwap:.3f}×{dynamic_factor:.2f}"
            
            # 添加交易时段信息
            if not is_trading_time:
                message += " [非交易时段]"
            
            return {
                'is_above': is_above,
                'vwap': vwap,
                'current_price': current_price,
                'factor': self.factor,
                'dynamic_factor': dynamic_factor,
                'message': message,
                'is_trading_time': is_trading_time
            }
            
        except Exception as e:
            return {
                'is_above': False,
                'vwap': None,
                'current_price': None,
                'factor': None,
                'dynamic_factor': None,
                'message': f'检查异常: {e}',
                'is_trading_time': False
            }
    
    def get_vwap_analysis(self, stock_code: str, end_time: Optional[datetime] = None) -> Dict[str, Union[float, int, str]]:
        """
        获取VWAP分析信息
        
        Args:
            stock_code: 股票代码
            end_time: 结束时间（回测模式使用，实时模式为None）
            
        Returns:
            包含分析信息的字典
        """
        vwap = self.get_latest_vwap(stock_code, end_time)
        if vwap is None:
            return {'error': '无法获取VWAP数据'}
        
        current_price = self._last_price.get(stock_code, 0)
        current_volume = self._last_volume.get(stock_code, 0)
        price_history = self._price_history.get(stock_code, [])
        volume_history = self._volume_history.get(stock_code, [])
        
        # 检查是否在交易时段
        if end_time is not None:
            # 回测模式：使用指定的结束时间
            now = end_time.time()
        else:
            # 实时模式：使用当前时间
            now = datetime.now().time()
            
        is_trading_time = ((self.MORNING_START <= now <= self.MORNING_END) or 
                          (self.AFTERNOON_START <= now <= self.AFTERNOON_END))
        
        # 计算价格相对于VWAP的偏离度
        price_deviation = ((current_price - vwap) / vwap * 100) if vwap > 0 else 0
        
        # 获取最近有效数据
        recent_prices = [p for p in price_history[-15:] if p > 0]  # 最近15个非0价格
        recent_volumes = [v for v in volume_history[-15:] if v > 0]  # 最近15个非0成交量
        
        return {
            'vwap': vwap,
            'current_price': current_price,
            'current_volume': current_volume,
            'price_deviation': price_deviation,
            'price_ratio': current_price / vwap if vwap > 0 and current_price > 0 else 0,
            'history_count': len(price_history),
            'valid_price_count': len([p for p in price_history if p > 0]),
            'valid_volume_count': len([v for v in volume_history if v > 0]),
            'min_price': min(price_history) if price_history else 0,
            'max_price': max(price_history) if price_history else 0,
            'avg_price': sum(price_history) / len(price_history) if price_history else 0,
            'min_volume': min(volume_history) if volume_history else 0,
            'max_volume': max(volume_history) if volume_history else 0,
            'avg_volume': sum(volume_history) / len(volume_history) if volume_history else 0,
            'is_trading_time': is_trading_time,
            'trading_status': '交易中' if is_trading_time else '休市中'
        }
    
    def force_refresh(self, stock_code: str) -> Optional[float]:
        """强制刷新指定股票的VWAP数据"""
        self.reset(stock_code)
        return self.get_latest_vwap(stock_code)
    
    def get_all_stock_vwap(self, stock_codes: list) -> Dict[str, Optional[float]]:
        """批量获取多个股票的VWAP值"""
        results = {}
        for code in stock_codes:
            results[code] = self.get_latest_vwap(code)
        return results


# 为了保持向后兼容，保留原有的类名
StandardVWAPCalculator = IndependentVWAPCalculator

# 更新导出列表
__all__ = ["IndependentVWAPCalculator", "StandardVWAPCalculator"]


# 使用示例
if __name__ == "__main__":
    # 定义股票代码
    stock_code = "513120.SH"  # 可以修改为您需要的股票代码
    
    # 创建VWAP计算器（周期为15，倍数为1.0）
    vwap_calc = IndependentVWAPCalculator(period=15, factor=1.0, xtdata=xtdata)
    target_time = datetime(2025, 8, 22, 14, 30, 1)  # 14:30:00
    
    # 获取最新VWAP
    vwap_value = vwap_calc.get_latest_vwap(stock_code, target_time)
    if vwap_value is not None:
        print(f"股票 {stock_code} 的最新VWAP: {vwap_value:.3f}")
        
        # 检查价格相对于VWAP的位置
        price_result = vwap_calc.check_price_vs_vwap(stock_code, target_time)
        print(f"价格判断: {price_result['message']}")
        
        # 获取详细分析
        analysis = vwap_calc.get_vwap_analysis(stock_code, target_time)
        if 'error' not in analysis:
            print(f"VWAP分析:")
            print(f"  当前价格: {analysis['current_price']:.3f}")
            print(f"  当前成交量: {analysis['current_volume']:.0f}")
            print(f"  价格偏离度: {analysis['price_deviation']:.2f}%")
            print(f"  价格比率: {analysis['price_ratio']:.3f}")
            print(f"  交易状态: {analysis['trading_status']}")
            print(f"  历史数据点: {analysis['history_count']} (价格有效: {analysis['valid_price_count']}, 成交量有效: {analysis['valid_volume_count']})")
        else:
            print(f"分析失败: {analysis['error']}")
    else:
        print(f"无法获取 {stock_code} 的VWAP值")
